"""Redis客户端 - 短期记忆存储"""

import json
import asyncio
from typing import Dict, Any, Optional, List, Union
from datetime import datetime, timedelta
from loguru import logger

try:
    import redis.asyncio as redis
except ImportError:
    import aioredis as redis

from src.config import get_database_url


class RedisClient:
    """Redis客户端 - 负责短期记忆存储"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.redis_config = config.get('databases', {}).get('redis', {})
        self.client: Optional[redis.Redis] = None
        self._connected = False
    
    async def connect(self):
        """连接Redis"""
        if self._connected:
            return
        
        try:
            redis_url = get_database_url(self.config, 'redis')
            self.client = redis.from_url(
                redis_url,
                encoding='utf-8',
                decode_responses=True,
                socket_timeout=self.redis_config.get('timeout', 5),
                socket_connect_timeout=self.redis_config.get('connect_timeout', 5),
                retry_on_timeout=True,
                health_check_interval=30
            )
            
            # 测试连接
            await self.client.ping()
            self._connected = True
            logger.info("Redis连接成功")
            
        except Exception as e:
            logger.error(f"Redis连接失败: {e}")
            raise
    
    async def close(self):
        """关闭Redis连接"""
        if self.client:
            await self.client.close()
            self._connected = False
            logger.info("Redis连接已关闭")
    
    async def ping(self) -> bool:
        """检查连接状态"""
        try:
            if self.client:
                await self.client.ping()
                return True
            return False
        except Exception:
            return False
    
    # 基础操作
    async def set(self, key: str, value: Any, expire: Optional[int] = None) -> bool:
        """设置键值对"""
        try:
            if isinstance(value, (dict, list)):
                value = json.dumps(value, ensure_ascii=False)
            
            result = await self.client.set(key, value, ex=expire)
            return result is True
        except Exception as e:
            logger.error(f"Redis SET操作失败 {key}: {e}")
            return False
    
    async def get(self, key: str, default: Any = None) -> Any:
        """获取值"""
        try:
            value = await self.client.get(key)
            if value is None:
                return default
            
            # 尝试解析JSON
            try:
                return json.loads(value)
            except (json.JSONDecodeError, TypeError):
                return value
        except Exception as e:
            logger.error(f"Redis GET操作失败 {key}: {e}")
            return default
    
    async def delete(self, *keys: str) -> int:
        """删除键"""
        try:
            return await self.client.delete(*keys)
        except Exception as e:
            logger.error(f"Redis DELETE操作失败 {keys}: {e}")
            return 0
    
    async def exists(self, key: str) -> bool:
        """检查键是否存在"""
        try:
            return await self.client.exists(key) > 0
        except Exception as e:
            logger.error(f"Redis EXISTS操作失败 {key}: {e}")
            return False
    
    async def expire(self, key: str, seconds: int) -> bool:
        """设置过期时间"""
        try:
            return await self.client.expire(key, seconds)
        except Exception as e:
            logger.error(f"Redis EXPIRE操作失败 {key}: {e}")
            return False
    
    async def ttl(self, key: str) -> int:
        """获取剩余生存时间"""
        try:
            return await self.client.ttl(key)
        except Exception as e:
            logger.error(f"Redis TTL操作失败 {key}: {e}")
            return -1
    
    # 列表操作
    async def lpush(self, key: str, *values: Any) -> int:
        """从左侧推入列表"""
        try:
            serialized_values = []
            for value in values:
                if isinstance(value, (dict, list)):
                    serialized_values.append(json.dumps(value, ensure_ascii=False))
                else:
                    serialized_values.append(str(value))
            
            return await self.client.lpush(key, *serialized_values)
        except Exception as e:
            logger.error(f"Redis LPUSH操作失败 {key}: {e}")
            return 0
    
    async def rpush(self, key: str, *values: Any) -> int:
        """从右侧推入列表"""
        try:
            serialized_values = []
            for value in values:
                if isinstance(value, (dict, list)):
                    serialized_values.append(json.dumps(value, ensure_ascii=False))
                else:
                    serialized_values.append(str(value))
            
            return await self.client.rpush(key, *serialized_values)
        except Exception as e:
            logger.error(f"Redis RPUSH操作失败 {key}: {e}")
            return 0
    
    async def lrange(self, key: str, start: int = 0, end: int = -1) -> List[Any]:
        """获取列表范围"""
        try:
            values = await self.client.lrange(key, start, end)
            result = []
            for value in values:
                try:
                    result.append(json.loads(value))
                except (json.JSONDecodeError, TypeError):
                    result.append(value)
            return result
        except Exception as e:
            logger.error(f"Redis LRANGE操作失败 {key}: {e}")
            return []
    
    async def llen(self, key: str) -> int:
        """获取列表长度"""
        try:
            return await self.client.llen(key)
        except Exception as e:
            logger.error(f"Redis LLEN操作失败 {key}: {e}")
            return 0
    
    # 哈希操作
    async def hset(self, key: str, field: str, value: Any) -> bool:
        """设置哈希字段"""
        try:
            if isinstance(value, (dict, list)):
                value = json.dumps(value, ensure_ascii=False)
            
            result = await self.client.hset(key, field, value)
            return result >= 0
        except Exception as e:
            logger.error(f"Redis HSET操作失败 {key}.{field}: {e}")
            return False
    
    async def hget(self, key: str, field: str, default: Any = None) -> Any:
        """获取哈希字段值"""
        try:
            value = await self.client.hget(key, field)
            if value is None:
                return default
            
            try:
                return json.loads(value)
            except (json.JSONDecodeError, TypeError):
                return value
        except Exception as e:
            logger.error(f"Redis HGET操作失败 {key}.{field}: {e}")
            return default
    
    async def hgetall(self, key: str) -> Dict[str, Any]:
        """获取所有哈希字段"""
        try:
            hash_data = await self.client.hgetall(key)
            result = {}
            for field, value in hash_data.items():
                try:
                    result[field] = json.loads(value)
                except (json.JSONDecodeError, TypeError):
                    result[field] = value
            return result
        except Exception as e:
            logger.error(f"Redis HGETALL操作失败 {key}: {e}")
            return {}
    
    async def get_hash(self, key: str) -> Optional[Dict[str, Any]]:
        """获取整个哈希对象"""
        try:
            hash_data = await self.client.hgetall(key)
            if not hash_data:
                return None
            
            # 尝试解析JSON值
            result = {}
            for k, v in hash_data.items():
                try:
                    # 尝试解析为JSON
                    result[k] = json.loads(v)
                except (json.JSONDecodeError, TypeError):
                    # 如果不是JSON，保持原始字符串
                    result[k] = v
            
            return result
            
        except Exception as e:
            logger.error(f"Redis GET_HASH操作失败 {key}: {e}")
            return None

    async def set_hash(self, key: str, value: Any, expire: Optional[int] = None) -> bool:
        """设置整个哈希对象"""
        try:
            if isinstance(value, dict):
                # 序列化字典中的复杂值
                serialized_value = {}
                for k, v in value.items():
                    if isinstance(v, (dict, list)):
                        serialized_value[k] = json.dumps(v, ensure_ascii=False)
                    else:
                        serialized_value[k] = str(v)
                
                # 使用hmset设置整个哈希
                result = await self.client.hset(key, mapping=serialized_value)
                
                # 设置过期时间
                if expire:
                    await self.expire(key, expire)
                
                return result >= 0
            else:
                logger.error(f"set_hash需要字典类型的值，得到: {type(value)}")
                return False
        except Exception as e:
            logger.error(f"Redis SET_HASH操作失败 {key}: {e}")
            return False
    
    # 短期记忆专用方法
    async def store_conversation_context(self, session_id: str, context: Dict[str, Any], ttl: int = 3600) -> bool:
        """存储对话上下文"""
        key = f"context:session:{session_id}"
        context['timestamp'] = datetime.now().isoformat()
        return await self.set(key, context, expire=ttl)
    
    async def get_conversation_context(self, session_id: str) -> Optional[Dict[str, Any]]:
        """获取对话上下文"""
        key = f"context:session:{session_id}"
        return await self.get(key)
    
    async def store_user_state(self, user_id: str, state: Dict[str, Any], ttl: int = 1800) -> bool:
        """存储用户状态"""
        key = f"state:user:{user_id}"
        state['last_updated'] = datetime.now().isoformat()
        return await self.set(key, state, expire=ttl)
    
    async def get_user_state(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户状态"""
        key = f"state:user:{user_id}"
        return await self.get(key)
    
    # 新增：短期记忆通用方法，供 MemoryManager 使用
    async def store_memory(self, user_id: str, memory_id: str, memory_data: Dict[str, Any], ttl: int = 3600) -> bool:
        """存储短期记忆到Redis，并维护用户最近记忆列表"""
        try:
            # 规范化时间字段，避免json序列化失败
            data = dict(memory_data)
            for ts_key in ("created_at", "last_accessed", "timestamp"):
                if ts_key in data and isinstance(data[ts_key], datetime):
                    data[ts_key] = data[ts_key].isoformat()

            data_key = f"memory:user:{user_id}:{memory_id}"
            recent_key = f"memories:recent:{user_id}"

            # 写入单条记忆数据（设置TTL）
            ok = await self.set(data_key, data, expire=ttl)
            if not ok:
                return False

            # 追加到最近列表（列表元素会被自动序列化）
            await self.lpush(recent_key, data)
            # 可选：设置最近列表的过期时间，避免无限增长
            await self.expire(recent_key, max(ttl, 3600))
            return True
        except Exception as e:
            logger.error(f"存储短期记忆失败 user={user_id} memory={memory_id}: {e}")
            return False

    async def get_memory(self, user_id: str, memory_id: str) -> Optional[Dict[str, Any]]:
        """获取短期记忆（如果存在）"""
        try:
            data_key = f"memory:user:{user_id}:{memory_id}"
            return await self.get(data_key)
        except Exception as e:
            logger.error(f"获取短期记忆失败 user={user_id} memory={memory_id}: {e}")
            return None

    async def extend_memory_ttl(self, user_id: str, memory_id: str, seconds: int) -> bool:
        """延长短期记忆的TTL"""
        try:
            data_key = f"memory:user:{user_id}:{memory_id}"
            return await self.expire(data_key, seconds)
        except Exception as e:
            logger.error(f"延长短期记忆TTL失败 user={user_id} memory={memory_id}: {e}")
            return False

    async def get_recent_memories(self, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取用户最近的短期记忆（从列表中读取）"""
        try:
            recent_key = f"memories:recent:{user_id}"
            return await self.lrange(recent_key, 0, max(0, limit - 1))
        except Exception as e:
            logger.error(f"获取最近短期记忆失败 user={user_id}: {e}")
            return []

    async def delete_memory(self, user_id: str, memory_id: str) -> bool:
        """删除短期记忆。注意：不会强制清理最近列表中的历史条目（非关键路径）。"""
        try:
            data_key = f"memory:user:{user_id}:{memory_id}"
            deleted = await self.delete(data_key)
            return deleted > 0
        except Exception as e:
            logger.error(f"删除短期记忆失败 user={user_id} memory={memory_id}: {e}")
            return False

    async def update_user_state(self, user_id: str, updates: Dict[str, Any], ttl: int = 1800) -> bool:
        """合并更新用户状态（短期）"""
        try:
            current = await self.get_user_state(user_id) or {}
            if not isinstance(current, dict):
                current = {}
            # 合并
            merged = {**current, **updates}
            merged['last_updated'] = datetime.now().isoformat()
            return await self.store_user_state(user_id, merged, ttl=ttl)
        except Exception as e:
            logger.error(f"更新用户状态失败 user={user_id}: {e}")
            return False

    async def get_stats(self) -> Dict[str, Any]:
        """返回Redis层的基础统计信息"""
        stats: Dict[str, Any] = {}
        try:
            stats['connected'] = await self.ping()
            try:
                stats['total_keys'] = await self.client.dbsize() if self.client else None
            except Exception:
                stats['total_keys'] = None

            # 统计与记忆相关的键数量
            try:
                mem_data_keys = await self.client.keys("memory:user:*") if self.client else []
                mem_recent_keys = await self.client.keys("memories:recent:*") if self.client else []
                stats['memory_keys'] = {
                    'data_keys': len(mem_data_keys),
                    'recent_lists': len(mem_recent_keys)
                }
            except Exception:
                stats['memory_keys'] = {
                    'data_keys': None,
                    'recent_lists': None
                }
        except Exception as e:
            logger.error(f"获取Redis统计信息失败: {e}")
            stats['error'] = str(e)
        return stats

    async def add_recent_interaction(self, session_id: str, interaction: Dict[str, Any], max_size: int = 50) -> bool:
        """添加最近交互记录"""
        key = f"interactions:recent:{session_id}"
        interaction['timestamp'] = datetime.now().isoformat()
        
        try:
            # 添加到列表左侧
            await self.lpush(key, interaction)
            
            # 保持列表大小
            current_size = await self.llen(key)
            if current_size > max_size:
                await self.client.ltrim(key, 0, max_size - 1)
            
            # 设置过期时间
            await self.expire(key, 7200)  # 2小时
            return True
        except Exception as e:
            logger.error(f"添加最近交互失败 {session_id}: {e}")
            return False
    
    async def get_recent_interactions(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取最近交互记录"""
        key = f"interactions:recent:{session_id}"
        return await self.lrange(key, 0, limit - 1)
    
    async def cache_search_result(self, query_hash: str, result: Any, ttl: int = 300) -> bool:
        """缓存搜索结果"""
        key = f"cache:search:{query_hash}"
        return await self.set(key, result, expire=ttl)
    
    async def get_cached_search_result(self, query_hash: str) -> Any:
        """获取缓存的搜索结果"""
        key = f"cache:search:{query_hash}"
        return await self.get(key)
    
    async def get_keys_pattern(self, pattern: str) -> List[str]:
        """根据模式获取键列表"""
        try:
            keys = await self.client.keys(pattern)
            return [key.decode() if isinstance(key, bytes) else key for key in keys]
        except Exception as e:
            logger.error(f"Redis KEYS操作失败 {pattern}: {e}")
            return []
    
    async def list_push(self, key: str, value: Any, max_length: Optional[int] = None) -> bool:
        """推入列表（从左侧），可选择限制列表长度"""
        try:
            # 序列化值
            if isinstance(value, (dict, list)):
                serialized_value = json.dumps(value, ensure_ascii=False)
            else:
                serialized_value = str(value)
            
            # 推入列表
            await self.client.lpush(key, serialized_value)
            
            # 如果设置了最大长度，裁剪列表
            if max_length and max_length > 0:
                await self.client.ltrim(key, 0, max_length - 1)
            
            return True
        except Exception as e:
            logger.error(f"Redis LIST_PUSH操作失败 {key}: {e}")
            return False
    
    async def list_range(self, key: str, start: int = 0, end: int = -1) -> List[Any]:
        """获取列表范围（别名方法，调用lrange）"""
        return await self.lrange(key, start, end)
    
    async def cleanup_expired_data(self):
        """清理过期数据"""
        try:
            # 获取所有键
            keys = await self.client.keys("*")
            expired_count = 0
            
            for key in keys:
                ttl = await self.ttl(key)
                if ttl == -2:  # 键已过期
                    expired_count += 1
            
            logger.info(f"清理了 {expired_count} 个过期键")
        except Exception as e:
            logger.error(f"清理过期数据失败: {e}")